package at.jku.ase.utils;

import at.jku.ase.encryption.Keys;

import javax.crypto.Cipher;
import javax.crypto.KeyGenerator;
import javax.crypto.SecretKey;
import javax.crypto.spec.SecretKeySpec;
import java.io.*;
import java.security.*;

public final class CryptUtils {

    private CryptUtils() {
    }

    public static final Keys<PublicKey, PrivateKey> generateRSAKeyPair() {
        return new OpWithParameterThrowsException<Keys<PublicKey, PrivateKey>>() {
            @Override
            public Keys<PublicKey, PrivateKey> op() throws Exception {
                KeyPairGenerator keyGen = KeyPairGenerator.getInstance(ASYMMETRIC_CIPHER_INSTANCE);
                SecureRandom random = new SecureRandom();
                keyGen.initialize(1024, random);
                KeyPair keyPair = keyGen.generateKeyPair();
                return new Keys<PublicKey, PrivateKey>(keyPair.getPublic(), keyPair.getPrivate());
            }
        }.execute();
    }

    public static final SecretKey generateAESKey() {
        return new OpWithParameterThrowsException<SecretKey>() {
            @Override
            public SecretKey op() throws Exception {
                KeyGenerator keyGen = KeyGenerator.getInstance(SYMMETRIC_CIPHER_INSTANCE);
                keyGen.init(128);
                return keyGen.generateKey();
            }
        }.execute();
    }

    private static final String SYMMETRIC_CIPHER_INSTANCE = "AES";
    private static final String ASYMMETRIC_CIPHER_INSTANCE = "RSA";

    public static final byte[] encryptSecretKey(final PublicKey publicKey, final SecretKey secretKey) {
        return new OpWithParameterThrowsException<byte[]>() {
            @Override
            public byte[] op() throws Exception {
                return asyCryptData(Cipher.ENCRYPT_MODE, secretKey.getEncoded(), publicKey);
            }
        }.execute();
    }

    public static final SecretKey decryptSecretKey(final PrivateKey privateKey, final byte[] keyCipher) {
        return new OpWithParameterThrowsException<SecretKey>() {
            @Override
            public SecretKey op() throws Exception {
                byte[] decryptedKeyBytes = asyCryptData(Cipher.DECRYPT_MODE, keyCipher, privateKey);
                return new SecretKeySpec(decryptedKeyBytes, SYMMETRIC_CIPHER_INSTANCE);
            }
        }.execute();
    }

    private static byte[] asyCryptData(final int cryptMode, final byte[] encoded, final Key key) {
        return new OpWithParameterThrowsException<byte[]>() {
            @Override
            public byte[] op() throws Exception {
                Cipher cipher = Cipher.getInstance("RSA"); //  RSA/ECB/PKCS1Padding
                cipher.init(cryptMode, key);
                return cryptData(encoded, cipher);
//                ByteArrayInputStream bis = new ByteArrayInputStream(encoded);
//                ByteArrayOutputStream bos = new ByteArrayOutputStream();
//                int blockSize = (1024 / 8) - 11;
//
//                while (bis.available() > 0) {
//                    byte[] buffer = new byte[Math.min(blockSize, bis.available())];
//                    bis.read(buffer);
//                    byte[] cipherBuffer = null;
//                    if (blockSize < bis.available()) {
//                        cipherBuffer = cipher.update(buffer);
//                    } else {
//                        cipherBuffer = cipher.doFinal(buffer);
//                    }
//                    bos.write(cipherBuffer);
//                }
//                return bos.toByteArray();
            }
        }.execute();
    }

    public static byte[] encryptData(final byte[] bytes, final SecretKey secretKey) {
        return new OpWithParameterThrowsException<byte[]>() {
            @Override
            public byte[] op() throws Exception {
                return symCryptData(Cipher.ENCRYPT_MODE, bytes, secretKey);
            }
        }.execute();
    }

    public static byte[] decryptData(final byte[] bytes, final SecretKey secretKey) {
        return new OpWithParameterThrowsException<byte[]>() {
            @Override
            public byte[] op() throws Exception {
                return symCryptData(Cipher.DECRYPT_MODE, bytes, secretKey);
            }
        }.execute();
    }

    private static byte[] symCryptData(final int cryptMode, final byte[] bytes, final SecretKey secretKey) {
        return new OpWithParameterThrowsException<byte[]>() {
            @Override
            public byte[] op() throws Exception {
                SecretKeySpec secretKeySpec = new SecretKeySpec(secretKey.getEncoded(), "AES");
                Cipher cipher = Cipher.getInstance("AES");
                cipher.init(cryptMode, secretKeySpec);
                return cryptData(bytes, cipher);
            }
        }.execute();
    }

    private static byte[] cryptData(final byte[] bytes, final Cipher cipher) {
        return new OpWithParameterThrowsException<byte[]>() {
            @Override
            public byte[] op() throws Exception {
                ByteArrayInputStream in = new ByteArrayInputStream(bytes);
                ByteArrayOutputStream out = new ByteArrayOutputStream();
                crypt(in, out, cipher);
                byte[] content = out.toByteArray();
                in.close();
                out.close();
                return content;
            }
        }.execute();
    }

    private static void crypt(InputStream in, OutputStream out, Cipher cipher) throws IOException,
            GeneralSecurityException {
        int blockSize = 1024/8;
        int outputSize = cipher.getOutputSize(blockSize);
        byte[] inBytes = new byte[blockSize];
        byte[] outBytes = new byte[outputSize];

        int inLength = 0;
        boolean more = true;
        while (more) {
            inLength = in.read(inBytes);
            if (inLength == blockSize) {
                int outLength = cipher.update(inBytes, 0, blockSize, outBytes);
                out.write(outBytes, 0, outLength);
            } else more = false;
        }
        if (inLength > 0) outBytes = cipher.doFinal(inBytes, 0, inLength);
        else outBytes = cipher.doFinal();
        out.write(outBytes);
    }

    public static byte[] signData(final byte[] data, final PrivateKey myPrivateKey) {
        return new OpWithParameterThrowsException<byte[]>() {
            @Override
            public byte[] op() throws Exception {
                Signature signature = Signature.getInstance("SHA1withRSA");
                signature.initSign(myPrivateKey);
                signature.update(data);
                return signature.sign();
            }
        }.execute();
    }


    public static boolean verifySignature(final byte[] data, final byte[] sign, final PublicKey publicKey) {
        return new OpWithParameterThrowsException<Boolean>() {
            @Override
            public Boolean op() throws Exception {
                Signature signature = Signature.getInstance("SHA1withRSA");
                signature.initVerify(publicKey);
                signature.update(data);
                return signature.verify(sign);
            }
        }.execute();
    }
}
